package com.lizba.redis.limit.slidinwindow;

import com.google.common.collect.ImmutableList;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.Pipeline;
import redis.clients.jedis.Response;

/**
 * <p>
 *     通过zset实现滑动窗口算法限流
 * </p>
 *
 * @Author: Liziba
 * @Date: 2021/9/6 18:11
 */
public class SimpleSlidingWindowByZSet {

    private Jedis jedis;

    public SimpleSlidingWindowByZSet(Jedis jedis) {
        this.jedis = jedis;
    }

    /**
     * 判断行为是否被允许
     *
     * @param userId        用户id
     * @param actionKey     行为key
     * @param period        限流周期（滑动窗口大小）
     * @param maxCount      最大请求次数
     * @return
     */
    public boolean isActionAllowed(String userId, String actionKey, int period, int maxCount) {
        String key = this.key(userId, actionKey);
        long ts = System.currentTimeMillis();
        Pipeline pipe = jedis.pipelined();
        pipe.multi();
        pipe.zadd(key, ts, String.valueOf(ts));
        // 移除滑动窗口之外的数据
        pipe.zremrangeByScore(key, 0, ts - (period * 1000));
        Response<Long> count = pipe.zcard(key);
        // 设置行为的过期时间，如果数据为冷数据，zset将会删除以此节省内存空间
        pipe.expire(key, period);
        pipe.exec();
        pipe.close();
        return count.get() <= maxCount;
    }


    /**
     * lua脚本限流
     *
     * @param userId
     * @param actionKey
     * @param period
     * @param maxCount
     * @return
     */
    public boolean isActionAllowedByLua(String userId, String actionKey, int period, int maxCount) {
        String luaScript = this.buildLuaScript();

        String key = key(userId, actionKey);
        long ts = System.currentTimeMillis();
        System.out.println(ts);
        ImmutableList<String> keys = ImmutableList.of(key);
        ImmutableList<String> args = ImmutableList.of(String.valueOf(ts),String.valueOf((ts - period * 1000)), String.valueOf(period));
        Number count = (Number) jedis.eval(luaScript, keys, args);

        return count != null && count.intValue() <= maxCount;
    }


    /**
     * 限流key
     *
     * @param userId
     * @param actionKey
     * @return
     */
    private String key(String userId, String actionKey) {
        return String.format("limit:%s:%s", userId, actionKey);
    }


    /**
     * 针对某个key使用lua脚本限流
     *
     * @return
     */
    private String buildLuaScript() {
        return "redis.call('ZADD', KEYS[1], tonumber(ARGV[1]), ARGV[1])" +
                "\nlocal c" +
                "\nc = redis.call('ZCARD', KEYS[1])" +
                "\nredis.call('ZREMRANGEBYSCORE', KEYS[1], 0, tonumber(ARGV[2]))" +
                "\nredis.call('EXPIRE', KEYS[1], tonumber(ARGV[3]))" +
                "\nreturn c;";

    }

}
